package com.gy.flink.state;

import org.apache.flink.api.common.functions.AggregateFunction;
import org.apache.flink.api.common.functions.ReduceFunction;
import org.apache.flink.api.common.functions.RichFlatMapFunction;
import org.apache.flink.api.common.state.*;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.java.tuple.Tuple;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.shaded.curator.org.apache.curator.shaded.com.google.common.collect.Lists;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.KeyedStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.util.Collector;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.UUID;
import java.util.stream.Collectors;

public class KeyedStateMain {


    public static void main(String[] args) throws Exception {
        Configuration conf = new Configuration();
        StreamExecutionEnvironment env = StreamExecutionEnvironment.createLocalEnvironmentWithWebUI(conf);

        DataStreamSource<Tuple2<Long, Long>> dataStream = env.fromElements(
                Tuple2.of(1L, 3L),
                Tuple2.of(1L, 5L),
                Tuple2.of(1L, 7L),
                Tuple2.of(2L, 4L),
                Tuple2.of(2L, 2L),
                Tuple2.of(2L, 5L)
        );

        KeyedStream<Tuple2<Long, Long>, Tuple> keyByStream = dataStream.keyBy(0);

//        SingleOutputStreamOperator<Tuple2<Long, Double>> result = keyByStream.flatMap(valueSateByFlatMap());
//        SingleOutputStreamOperator<Tuple2<Long, Double>> result = keyByStream.flatMap(listSateByFlatMap());
//        SingleOutputStreamOperator<Tuple2<Long, Double>> result = keyByStream.flatMap(mapSateByFlatMap());
//        result.print().setParallelism(1);

//        keyByStream.sum(1).print();
//        keyByStream.flatMap(reducingSateByFlatMap()).print();
        keyByStream.flatMap(aggregatingStateSateByFlatMap()).print();


        env.execute(KeyedStateMain.class.getSimpleName());
    }

    private static RichFlatMapFunction<Tuple2<Long, Long>, Tuple2<Long, Double>> valueSateByFlatMap() {

        return new RichFlatMapFunction<Tuple2<Long, Long>, Tuple2<Long, Double>>() {

            /**
             * k:出现次数
             * v:和
             */
            ValueState<Tuple2<Long, Long>> state;

            @Override
            public void open(Configuration parameters) throws Exception {
                //注册state 让flink管理state
                ValueStateDescriptor<Tuple2<Long, Long>> descriptor = new ValueStateDescriptor<>(
                        "avg",
                        Types.TUPLE(Types.LONG, Types.LONG)
                );
                state = super.getRuntimeContext().getState(descriptor);
            }

            @Override
            public void flatMap(Tuple2<Long, Long> element, Collector<Tuple2<Long, Double>> out) throws Exception {
                Tuple2<Long, Long> currStae = state.value();
                if (currStae == null) {
                    currStae = Tuple2.of(0L, 0L);
                }

                currStae.f0++;
                currStae.f1 += element.f1;

                state.update(currStae);

                if (currStae.f0 == 3) {
                    out.collect(Tuple2.of(element.f0, (double) currStae.f1 / currStae.f0));
                    //清除历史状态
                    state.clear();
                }
            }
        };
    }


    private static RichFlatMapFunction<Tuple2<Long, Long>, Tuple2<Long, Double>> listSateByFlatMap() {

        return new RichFlatMapFunction<Tuple2<Long, Long>, Tuple2<Long, Double>>() {

            /**
             * k:原始key
             * v:原始value
             */
            ListState<Tuple2<Long, Long>> state;

            @Override
            public void open(Configuration parameters) throws Exception {
                //注册state 让flink管理state
                ListStateDescriptor<Tuple2<Long, Long>> descriptor = new ListStateDescriptor<>(
                        "avg",
                        Types.TUPLE(Types.LONG, Types.LONG)
                );
                state = super.getRuntimeContext().getListState(descriptor);
            }

            @Override
            public void flatMap(Tuple2<Long, Long> element, Collector<Tuple2<Long, Double>> out) throws Exception {
                Iterable<Tuple2<Long, Long>> currStates = state.get();
                if (currStates == null) {
                    state.addAll(Collections.emptyList());
                }

                state.add(element);

                ArrayList<Tuple2<Long, Long>> arrayList = Lists.newArrayList(currStates);
                if (arrayList.size() == 3) {
                    out.collect(Tuple2.of(element.f0, arrayList.stream().limit(3L).collect(Collectors.averagingDouble(x -> x.f1))));
                    //清除历史状态
                    state.clear();
                }
            }
        };
    }


    private static RichFlatMapFunction<Tuple2<Long, Long>, Tuple2<Long, Double>> mapSateByFlatMap() {

        return new RichFlatMapFunction<Tuple2<Long, Long>, Tuple2<Long, Double>>() {

            /**
             * k:原始uuid_key
             * v:原始value
             */
            MapState<String, Long> state;

            @Override
            public void open(Configuration parameters) throws Exception {
                //注册state 让flink管理state
                MapStateDescriptor<String, Long> descriptor = new MapStateDescriptor(
                        "avg",
                        String.class, Long.class
                );
                state = super.getRuntimeContext().getMapState(descriptor);
            }

            @Override
            public void flatMap(Tuple2<Long, Long> element, Collector<Tuple2<Long, Double>> out) throws Exception {
                state.put(UUID.randomUUID().toString() + "_" + element.f0, element.f1);

                ArrayList<String> keyList = Lists.newArrayList(state.keys());

                if (keyList.size() == 3) {
                    Double avg = keyList.stream().limit(3L).collect(Collectors.averagingDouble(x -> {
                        Long aLong = null;
                        try {
                            aLong = state.get(x);
                        } catch (Exception e) {
                            e.printStackTrace();
                        }
                        return aLong;
                    }));
                    out.collect(Tuple2.of(element.f0, avg));

                    //清除历史状态
                    state.clear();
                }
            }
        };
    }


    private static RichFlatMapFunction<Tuple2<Long, Long>, Tuple2<Long, Long>> reducingSateByFlatMap() {

        return new RichFlatMapFunction<Tuple2<Long, Long>, Tuple2<Long, Long>>() {

            /**
             * 保存同一个key的值
             */
            ReducingState<Long> state;

            @Override
            public void open(Configuration parameters) throws Exception {
                //注册state 让flink管理state
                ReducingStateDescriptor<Long> descriptor = new ReducingStateDescriptor<Long>(
                        "sum",
                        new ReduceFunction<Long>() {
                            @Override
                            public Long reduce(Long v1, Long v2) throws Exception {
                                return v1 + v2;
                            }
                        }, Long.class
                );
                state = super.getRuntimeContext().getReducingState(descriptor);
            }

            @Override
            public void flatMap(Tuple2<Long, Long> element, Collector<Tuple2<Long, Long>> out) throws Exception {
                state.add(element.f1);
                out.collect(Tuple2.of(element.f0, state.get()));
            }
        };
    }


    private static RichFlatMapFunction<Tuple2<Long, Long>, Tuple2<Long, String>> aggregatingStateSateByFlatMap() {

        return new RichFlatMapFunction<Tuple2<Long, Long>, Tuple2<Long, String>>() {

            /**
             * 保存同一个key的值
             */
            AggregatingState<Long, String> state;

            @Override
            public void open(Configuration parameters) throws Exception {
                //注册state 让flink管理state   输入类型，辅助类型，输出类型
                AggregatingStateDescriptor<Long, List<String>, String> descriptor = new AggregatingStateDescriptor<>("list_str", new AggregateFunction<Long, List<String>, String>() {

                    /**
                     * 创建辅助变量
                     * @return
                     */
                    @Override
                    public List<String> createAccumulator() {
                        ArrayList<String> initList = new ArrayList<>();
                        initList.add("Contains");
                        return initList;
                    }

                    /**
                     * 每次的操作
                     * @param currVal
                     * @param list
                     * @return
                     */
                    @Override
                    public List<String> add(Long currVal, List<String> list) {
                        list.add("add");
                        list.add(currVal.toString());
                        return list;
                    }

                    /**
                     * 合并task的数据
                     * @param currTaskList
                     * @param otherTaskList
                     * @return
                     */
                    @Override
                    public List<String> merge(List<String> currTaskList, List<String> otherTaskList) {
                        currTaskList.addAll(otherTaskList);
                        return currTaskList;
                    }

                    /**
                     * 最终结果
                     * @param resultList
                     * @return
                     */
                    @Override
                    public String getResult(List<String> resultList) {
                        return String.join(" ", resultList);
                    }
                }, Types.LIST(Types.STRING));//辅助类型

                state = super.getRuntimeContext().getAggregatingState(descriptor);
            }

            @Override
            public void flatMap(Tuple2<Long, Long> element, Collector<Tuple2<Long, String>> out) throws Exception {
                state.add(element.f1);
                out.collect(Tuple2.of(element.f0, state.get()));
            }
        };
    }

}
